import os
import random
import torch
import numpy as np
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
    BitsAndBytesConfig,
    set_seed
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

from config import output_dir, model_path, data_path, seed, max_length

# ========= 配置 =========

log_dir = os.path.join(output_dir, "logs")  # TensorBoard 日志路径

batch_size = 10
num_epochs = 2
learning_rate = 2e-5
lora_rank = 16
lora_alpha = 32
lora_dropout = 0.05
use_8bit = True


# ========= 固定随机性 =========
set_seed(seed)
random.seed(seed)
np.random.seed(seed)

# ========= 加载模型 & tokenizer =========
print(f"Loading model from {model_path} ...")
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True
)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    quantization_config=quant_config,
    trust_remote_code=True
)

# ========= LoRA 配置 =========
lora_config = LoraConfig(
    r=lora_rank,
    lora_alpha=lora_alpha,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=lora_dropout,
    bias="none",
    task_type="CAUSAL_LM"
)
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, lora_config)

# ========= 加载数据集 =========
print(f"Loading dataset from {data_path} ...")
dataset = load_dataset("json", data_files=data_path)["train"]

# ========= 划分训练 / 测试集 =========
split_dataset = dataset.train_test_split(test_size=0.1, seed=seed)
train_dataset = split_dataset["train"]
test_dataset = split_dataset["test"]
print(f"Train samples: {len(train_dataset)}, Test samples: {len(test_dataset)}")

# ========= 构造 prompt =========
def format_example(example):
    instruction = example["instruction"]
    input_text = example.get("input", "")
    output = example["output"]

    # 构造 prompt
    if input_text.strip():
        prompt = f"指令：{instruction}\n输入：{input_text}\n回答："
    else:
        prompt = f"指令：{instruction}\n回答："

    # 拼接完整输入（prompt + output）
    text = prompt + output

    # 分别计算 prompt 的长度
    prompt_tokens = tokenizer(prompt, truncation=True, max_length=max_length)
    prompt_len = len(prompt_tokens["input_ids"])

    # 对整个 text 做 tokenization
    tokenized = tokenizer(
        text,
        truncation=True,
        max_length=max_length,
        padding="max_length"
    )

    # 构造 labels：prompt 部分标记为 -100，不计入 loss
    labels = tokenized["input_ids"].copy()
    labels[:prompt_len] = [-100] * prompt_len
    labels[labels == tokenizer.pad_token_id] = -100
    tokenized["labels"] = labels

    return tokenized


tokenized_train = train_dataset.map(format_example)
tokenized_test = test_dataset.map(format_example)

# ========= 训练配置 =========
training_args = TrainingArguments(
    output_dir=output_dir,
    logging_dir=log_dir,                # TensorBoard 日志目录
    per_device_train_batch_size=batch_size,
    gradient_accumulation_steps=8,
    learning_rate=learning_rate,
    num_train_epochs=num_epochs,
    lr_scheduler_type="cosine",
    warmup_ratio=0.05,
    fp16=True,
    logging_steps=10,                   # 每10步记录一次loss
    save_strategy="epoch",
    save_total_limit=2,
    report_to=["tensorboard"],          # 启用TensorBoard
)

# ========= Trainer =========
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_test,
    data_collator=data_collator
)

# ========= 开始训练 =========
trainer.train()

# ========= 保存模型 =========
model.save_pretrained(os.path.join(output_dir, "lora"))
tokenizer.save_pretrained(os.path.join(output_dir, "lora"))
print("\n✅ 微调完成！模型已保存到：", os.path.join(output_dir, "lora"))
print(f"📊 TensorBoard 日志已保存到：{log_dir}")

# ========= 查看方式 =========
print("\n🚀 你可以在另一个终端运行以下命令实时查看loss曲线：")
print(f"tensorboard --logdir={log_dir} --port=6006")
print("然后在浏览器打开：http://localhost:6006/")
